import logging
import os
from contextlib import AsyncExitStack
from typing import Dict, List, Any, Optional, AsyncContextManager

# 导入官方的MCP Python SDK
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client
# 添加SSE相关导入
from mcp.client.sse import sse_client

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("mcp")


class MCPClient:
    """MCP客户端，用于连接和交互MCP服务器"""

    def __init__(self, client_name: str = "Wend MCP Client", client_version: str = "1.0.0"):
        """初始化MCP客户端

        Args:
            client_name: 客户端名称
            client_version: 客户端版本
        """
        self.client_name = client_name
        self.client_version = client_version
        self.session = None
        self.exit_stack = AsyncExitStack()
        self.stdio = None
        self.write = None
        self.connection_type = None  # 记录连接类型：stdio 或 sse
        logger.info(f"创建MCP客户端: {client_name} v{client_version}")

    async def connect_with_sse(self, base_url: str, session_id: Optional[str] = None):
        """使用SSE连接到MCP服务器

        Args:
            base_url: MCP服务器基础URL，例如 "http://localhost:3000"
            session_id: 会话ID，用于区分不同的客户端连接

        Returns:
            self: 返回自身实例以支持链式调用
        """
        logger.info(f"使用SSE连接到MCP服务器: {base_url}")

        try:
            # 创建SSE客户端
            sse_transport = await self.exit_stack.enter_async_context(sse_client(url=base_url))
            
            # 创建会话并初始化
            # 使用正确的对象属性
            self.session = await self.exit_stack.enter_async_context(ClientSession(
                sse_transport[0],  # MemoryObjectReceiveStream
                sse_transport[1]   # MemoryObjectSendStream
            ))
            await self.session.initialize()
            
            # 标记连接类型
            self.connection_type = "sse"

            # 获取可用工具列表
            response = await self.session.list_tools()
            tools = response.tools
            logger.info(f"SSE连接成功，服务器提供了 {len(tools)} 个工具")
            return self
        except Exception as e:
            logger.error(f"SSE连接MCP服务器失败: {str(e)}")
            await self.close()
            raise

    async def connect_with_command(self, command: str, args: Optional[List[str]] = None, env: Optional[Dict[str, str]] = None):
        """使用命令和参数连接到MCP服务器

        Args:
            command: 命令
            args: 命令参数列表
            env: 环境变量
        """
        if args is None:
            args = []
        if env is None:
            env = {}

        logger.info(f"连接到MCP服务器: {command} {' '.join(args)}")

        try:
            # 创建服务器参数
            server_params = StdioServerParameters(
                command=command,
                args=args,
                env=env
            )

            # 创建stdio客户端
            stdio_transport = await self.exit_stack.enter_async_context(stdio_client(server_params))
            self.stdio, self.write = stdio_transport

            # 创建会话并初始化
            self.session = await self.exit_stack.enter_async_context(ClientSession(self.stdio, self.write))
            await self.session.initialize()
            
            # 标记连接类型
            self.connection_type = "stdio"

            # 获取可用工具列表
            response = await self.session.list_tools()
            tools = response.tools
            logger.info(f"连接成功，服务器提供了 {len(tools)} 个工具")
            return self
        except Exception as e:
            logger.error(f"连接MCP服务器失败: {str(e)}")
            await self.close()
            raise

    async def connect_with_script(self, script_path: str):
        """连接到Python或JavaScript脚本MCP服务器

        Args:
            script_path: 脚本路径
        """
        is_python = script_path.endswith('.py')
        is_js = script_path.endswith('.js')

        if not (is_python or is_js):
            raise ValueError("服务器脚本必须是.py或.js文件")

        # 确保脚本路径是绝对路径
        if not os.path.isabs(script_path):
            script_path = os.path.abspath(script_path)

        # 根据脚本类型选择适当的命令
        if is_python:
            # 检查是否有uvicorn
            try:
                import importlib.util
                if importlib.util.find_spec("uvicorn"):
                    command = "uvicorn"
                    # 获取模块名，移除.py后缀
                    module_name = os.path.basename(script_path).replace('.py', '')
                    return await self.connect_with_command(command, [f"{module_name}:app", "--reload"])
            except ImportError:
                pass

            # 回退到直接使用python
            command = "python"
        else:
            command = "node"

        return await self.connect_with_command(command, [script_path])

    async def list_tools(self):
        """列出服务器上的所有工具

        Returns:
            工具列表
        """
        self._ensure_connected()
        response = await self.session.list_tools()
        return response.tools

    async def call_tool(self, name: str, arguments: Dict[str, Any]):
        """调用指定的工具

        Args:
            name: 工具名称
            arguments: 工具参数

        Returns:
            工具调用结果
        """
        self._ensure_connected()
        try:
            response = await self.session.call_tool(name, arguments=arguments)
            return response
        except Exception as e:
            logger.error(f"调用工具 {name} 失败: {str(e)}")
            raise

    async def list_resources(self):
        """列出服务器上的所有资源

        Returns:
            资源列表
        """
        self._ensure_connected()
        response = await self.session.list_resources()
        return response.resources

    async def read_resource(self, uri: str):
        """读取指定URI的资源

        Args:
            uri: 资源URI

        Returns:
            资源内容
        """
        self._ensure_connected()
        response = await self.session.read_resource(uri)
        return response.content

    async def list_prompts(self):
        """列出服务器上的所有提示

        Returns:
            提示列表
        """
        self._ensure_connected()
        response = await self.session.list_prompts()
        return response.prompts

    async def execute_prompt(self, name: str, arguments: Optional[Dict[str, Any]] = None):
        """执行指定的提示模板

        Args:
            name: 提示名称
            arguments: 提示参数

        Returns:
            提示执行结果
        """
        if arguments is None:
            arguments = {}

        self._ensure_connected()
        response = await self.session.execute_prompt(name, arguments=arguments)
        return response

    async def get_prompt(self, name: str, arguments: Optional[Dict[str, Any]] = None):
        """获取指定的提示内容

        Args:
            name: 提示名称
            arguments: 提示参数

        Returns:
            提示内容
        """
        if arguments is None:
            arguments = {}

        self._ensure_connected()
        response = await self.session.get_prompt(name, arguments=arguments)
        return response

    async def ping(self):
        """发送ping请求到服务器

        Returns:
            服务器响应
        """
        self._ensure_connected()
        response = await self.session.ping()
        return response

    def _ensure_connected(self):
        """确保客户端已连接到服务器"""
        if self.session is None:
            raise RuntimeError("客户端未连接到服务器，请先调用connect_with_command或connect_with_script方法")

    async def close(self):
        """关闭客户端连接"""
        if self.session is not None:
            await self.exit_stack.aclose()
            self.session = None
            self.stdio = None
            self.write = None
            self.connection_type = None
            logger.info("已关闭MCP客户端连接")


class MCPClientContext(AsyncContextManager):
    """MCP客户端上下文管理器，支持async with语句"""

    def __init__(self, tool_config: Optional[Dict[str, Any]] = None, command: Optional[str] = None, args: Optional[List[str]] = None, 
                 base_url: Optional[str] = None, session_id: Optional[str] = None):
        """初始化MCP客户端上下文

        Args:
            tool_config: 工具配置字典
            command: 命令（当不使用tool_config时）
            args: 命令参数（当不使用tool_config时）
            base_url: SSE服务器URL（当使用SSE连接时）
            session_id: SSE会话ID（当使用SSE连接时）
        """
        self.tool_config = tool_config
        self.command = command
        self.args = args
        self.base_url = base_url
        self.session_id = session_id
        self.client = None

    async def __aenter__(self) -> MCPClient:
        """进入上下文，连接到服务器"""
        self.client = MCPClient()

        try:
            if self.tool_config:
                # 从工具配置连接
                connection_type = self.tool_config.get("connection_type", "stdio")
                
                if connection_type == "sse":
                    # 使用SSE连接
                    base_url = self.tool_config.get("base_url")
                    session_id = self.tool_config.get("session_id")
                    if not base_url:
                        raise ValueError("使用SSE连接时必须提供base_url")
                    return await self.client.connect_with_sse(base_url, session_id)
                else:
                    # 使用stdio连接
                    command = self.tool_config.get("command")
                    args = self.tool_config.get("args", [])
                    env = self.tool_config.get("env", {})

                    # 如果tool_config指定了server_script，优先使用它
                    server_script = self.tool_config.get("server_script")
                    if server_script and not args:
                        return await self.client.connect_with_script(server_script)

                    return await self.client.connect_with_command(command, args, env)
            elif self.base_url:
                # 使用SSE连接
                return await self.client.connect_with_sse(self.base_url, self.session_id)
            elif self.command:
                # 使用命令连接
                return await self.client.connect_with_command(self.command, self.args or [])
            else:
                raise ValueError("必须提供工具配置、命令或SSE服务器URL")

        except Exception as e:
            # 如果连接失败，确保关闭客户端
            if self.client:
                await self.client.close()
            raise e

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """退出上下文，关闭连接"""
        if self.client:
            await self.client.close()